import torch
from   xllm import _custom_ops


a = torch.rand((2,3), dtype=torch.float32)
b = torch.rand((2,3), dtype=torch.float32)
c = torch.zeros_like(a)
val = 2.5

############################################################
print("\n\n++++++++++++ CPU后端验证++++++++++++++")
print("\n=====\nmuladd: ")
result = _custom_ops.muladd(a, b, val)
print(torch.allclose(result, a*b+val))

print("\n=====\nmul: ")
result = _custom_ops.mul(a, b)
print(torch.allclose(result, a*b))

print("\n=====\nadd: ")
c.fill_(0)
_custom_ops.add(a, b, c)
print(torch.allclose(c, a+b))



############################################################
print("\n\n++++++++++++ CUDA后端验证++++++++++++++")
a = a.cuda()
b = b.cuda()
c = c.cuda()
val = 2.5
print("\n=====\nmuladd: ")
result = _custom_ops.muladd(a, b, val)
print(torch.allclose(result, a*b+val))

print("\n=====\nmul: ")
result = _custom_ops.mul(a, b)
print(torch.allclose(result, a*b))

print("\n=====\nadd: ")
c.fill_(0)
_custom_ops.add(a, b, c)
print(torch.allclose(c, a+b))

############################################################
print("\n\n++++++++++++ torch.compile验证++++++++++++++")
@torch.compile(fullgraph=True)
def auto_graph(a:torch.Tensor, b:torch.Tensor):
    return _custom_ops.mul(a, b)

print(torch.allclose(auto_graph(a, b), a*b))